#!/usr/bin/env python
#
# Copyright (C) 2007 Jason Kivlighn
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, #USA.

import time
import sys, os
import random
import cairo
import sqlite3

import timeline_data as tl
import mtzcolor

class Graph:
  WEEK = 7*24*60*60 # seconds
  def __init__(self,data,options):
    self.data = data
    
    self.define_defaults()
    for key, value in options.iteritems():
      self.__dict__[key] = value

  def define_defaults(self):
    self.times = None
    self.artist_stack = None
    self.color_stack = None
    self.user = None
    self.x_factor = 40
    self.y_factor = 1.0
    self.fill = True
    self.min_height = 300
    self.min_font_size = 6
    self.max_font_size = 27
    self.tick_interval = 3
    
    self.x_offset = 0
    self.y_offset = 30
    
    self.black_on_white = True
    if self.black_on_white:
      self.bg_color = (1,1,1)
      self.x_axis_color = (0,0,0)
      self.title_color = (0,0,0)
      self.label_color = (0,0,0)
      self.x_grid_color = (0,0,0)
    else:
      self.bg_color = (0.0,0.0,0.0)
      self.x_axis_color = (1,1,1)
      self.title_color = (1,1,1)
      self.label_color = (1,1,1)
      self.x_grid_color = (1,1,1)
    
    self.font = "Sans"
    self.hide_x_axis = False
    self.hide_y_axis = True
    self.x_axis_label_angle = -50.0

  def draw(self,context):
    context.save()
    context.translate(self.x_offset,self.y_offset)
    
    for i in range(1,len(self.data[0])):
      print "%.2f%%" % (100.0*i/len(self.data[0]))
      line = [[d[0],d[-i]] for d in self.data]
      context.move_to(0,0)
      for j in range(len(line)):
        context.line_to(line[j][0]*self.x_factor,line[j][1]*self.y_factor)
      context.line_to(line[-1][0]*self.x_factor,0)
      if i < len(self.data[0])-1:
        r,g,b = self.color_stack[self.artist_stack[i-1]]
      else:
        r,g,b = self.bg_color
      context.set_source_rgb(r, g, b)
      if self.fill:
        context.fill()
      else:
        context.stroke()
      context.close_path()
    
    context.set_source_rgb(*self.label_color)
    context.select_font_face(self.font,cairo.FONT_SLANT_NORMAL,cairo.FONT_WEIGHT_NORMAL)
    
    self.draw_labels(context)
    if not self.hide_x_axis:
      self.draw_x_axis(context)
    context.restore()
    self.draw_title(context)
    
    if not self.hide_y_axis:
      self.draw_y_axis(context)

  def label_size(self,text,context):
    return context.text_extents(text)[2:4]

  def draw_labels(self,context):
    label_positions = tl.label_positions_where_popular(self.data,self.artist_stack,
      lambda text: self.label_size(text,context),
      self.min_font_size,self.max_font_size)
    for artist,x,y,size in label_positions:
      context.set_font_size(size)
      width, height = self.label_size(artist,context)
      context.move_to(min(self.data[-1][0]*self.x_factor-width-1,
         max(1,
          x*self.x_factor-width/2.0))
        ,y*self.y_factor-height/2.0)
      context.show_text(artist)
  
  def draw_title(self,context):
    context.set_source_rgb(*self.title_color)
    context.select_font_face(self.font,cairo.FONT_SLANT_NORMAL,cairo.FONT_WEIGHT_NORMAL)
    context.set_font_size(20)
    
    begin, end = (self.times[0],self.times[-1]+self.WEEK)
    text = "%s from %s to %s" % (self.user,time.strftime("%b %Y",time.gmtime(begin)),time.strftime("%b %Y",time.gmtime(end)))
    text_width, text_height = context.text_extents(text)[2:4]
    
    context.move_to(self.data[len(self.data)/2][0]*self.x_factor-text_width/2,text_height)
    context.show_text(text)

  def draw_x_axis(self,context):
    max_max = max([max(l[1:]) for l in self.data])
    
    context.select_font_face(self.font,cairo.FONT_SLANT_NORMAL,cairo.FONT_WEIGHT_BOLD)
    context.set_font_size(12)
    
    context.set_dash([2])
    context.set_line_width(0.3)
    
    for i in range(0,len(self.times),self.tick_interval):
      text = time.strftime("%b %d, %Y",time.gmtime(self.times[i]))
      text_width, text_height = context.text_extents(text)[2:4]
    
      context.set_source_rgb(*self.x_axis_color)
      
      context.save()
      context.move_to(i*self.x_factor,max_max*self.y_factor)
      context.rotate(self.x_axis_label_angle)
      context.show_text(text)
      context.restore();
      
      context.move_to(i*self.x_factor,0)
      context.line_to(i*self.x_factor,max_max*self.y_factor)
      context.set_source_rgb(*self.x_grid_color) 
      context.stroke()
      
  def draw_y_axis(self,context):
    max_max = max([max(l[1:]) for l in self.data])
    center = max_max/2
    
    context.select_font_face(self.font,cairo.FONT_SLANT_NORMAL,cairo.FONT_WEIGHT_BOLD)
    context.set_font_size(12)
    
    context.set_dash([])
    context.set_line_width(1)
    
    context.move_to(self.x_offset,self.y_offset)
    context.line_to(self.x_offset,self.y_offset+max_max*self.y_factor)
    context.stroke()
    
    text = unicode(int(max_max))
    width, height = context.text_extents(text)[2:4]
    context.move_to(self.x_offset/2-width/2,self.y_offset+height/2)
    context.show_text(text)
    context.move_to(self.x_offset/2-width/2,self.y_offset+height/2+max_max*self.y_factor)
    context.show_text(text)
    
    text = unicode(0)
    width, height = context.text_extents(text)[2:4]
    context.move_to(self.x_offset/2-width/2,self.y_offset+height/2+center)
    context.show_text(text)

try:
  import gtk
except:
  sys.stderr.write("""WARNING: gtk module not available
  Unable to display graph.  Graph will be written to disk instead.""")
  have_gtk = False
else:
  have_gtk = True
  class GtkGraph (gtk.Image):
    def __init__(self,g,width,height):
      gtk.Image.__init__(self)
      
      colormap = gtk.gdk.colormap_get_system()
      pixmap = gtk.gdk.Pixmap(None,width,height,colormap.get_visual().depth)
      pixmap.set_colormap(colormap)
      self.set_from_pixmap(pixmap,None)
      context = pixmap.cairo_create()
      context.rectangle(0,0,width,height)
      context.set_source_rgb(*g.bg_color)
      context.fill()
      g.draw(context)
    
def display_graph(data,options):
  g = Graph(data,options)
  max_max = max([max(l[1:]) for l in data])
  height = max_max*g.y_factor+g.y_offset
  
  if g.min_height:
    if height < g.min_height:
      g.y_factor = g.min_height/max_max
      height = g.min_height+g.y_offset

  height = int(height)
  width = int(data[-1][0]*g.x_factor)
  
  image_out = options.get("output_image")
  if not image_out and not have_gtk:
    image_out = "%s.png" % options["user"]
  if image_out:
    surface = cairo.ImageSurface(cairo.FORMAT_ARGB32, width, height+30)
    #surface = cairo.SVGSurface(image_out, width, height)
    context = cairo.Context(surface)
    context.rectangle(0,0,width,height+30)
    context.set_source_rgb(*g.bg_color)
    context.fill()
    g.draw(context)
    if image_out == "-":
      surface.write_to_png(sys.stdout)
    else:
      print "Writing %s" % image_out
      surface.write_to_png(image_out)

  no_display = options.get("no_display")
  if not no_display and have_gtk:
    def destroy(window):
      gtk.main_quit()
    
    window = gtk.Window()
    window.set_default_size(800,700)
    swindow = gtk.ScrolledWindow()
    widget = GtkGraph(g,width,height+30)
    window.show()
    swindow.show()
    widget.show()
    swindow.add_with_viewport(widget)
    window.add(swindow)
    window.connect("destroy",destroy)
    gtk.main()

if __name__=="__main__":
  if len(sys.argv) < 2 or "--help" in sys.argv:
    print """Usage: ./cairo_timeline.py <user> [OPTIONS]
    
    OPTIONS:
      -type (artists|tracks|albums|track_tags|artist_tags)
      -threshold <X>
      -to-disk <filename.png>
      -no-display
      -style (raw|linregress|spline)
      -no-fill
      -random-colors
      -font-range <min>,<max>
      -font <name>
      -hide-x-axis
      -interval <num of weeks>
      -h-stretch-factor <X>
      -v-stretch-factor <X>
      -tick-interval <num of weeks>
      -no-stack
      """
    sys.exit(1)
  user = sys.argv[1]
  
  options = {}
  
  if "-type" in sys.argv:
    type = sys.argv[sys.argv.index("-type")+1]
  else:
    type = "artists"
    
  if "-threshold" in sys.argv:
    threshold = int(sys.argv[sys.argv.index("-threshold")+1])
  else:
    threshold = 5
    
  if "-to-disk" in sys.argv:
    image = sys.argv[sys.argv.index("-to-disk")+1]
  else:
    image = None

  if "-style" in sys.argv:
    style = sys.argv[sys.argv.index("-style")+1]
  else:
    style = "spline"
    
  if "-font-range" in sys.argv:
    min_font, max_font = map(int,sys.argv[sys.argv.index("-font-range")+1].split(","))
    options["min_font_size"] = min_font
    options["max_font_size"] = max_font
    
  if "-font" in sys.argv:
    font = sys.argv[sys.argv.index("-font")+1]
    options["font"] = font
    
  if "-h-stretch-factor" in sys.argv:
    f = float(sys.argv[sys.argv.index("-h-stretch-factor")+1])
    options["x_factor"] = f
    
  if "-v-stretch-factor" in sys.argv:
    f = float(sys.argv[sys.argv.index("-v-stretch-factor")+1])
    options["y_factor"] = f
    
  if "-interval" in sys.argv:
    interval = int(sys.argv[sys.argv.index("-interval")+1])
  else:
    interval = 1
    
  stack = "-no-stack" not in sys.argv
  centered = stack
    
  if "-tick-interval" in sys.argv:
    f = int(sys.argv[sys.argv.index("-tick-interval")+1])
    options["tick_interval"] = f
  
  db_name = "lastfm-%s.db" % user
  if not os.path.exists(db_name):
    print "User database '%s' doesn't exist" % user
    sys.exit(1)
  conn = sqlite3.connect("lastfm-%s.db" % user)
    
  if style == "raw":
    times, artist_stack, data = tl.raw_data(conn,type,user,threshold=threshold,interval=interval,stack=stack)
    if centered:
      tl.center_data(data)
  elif style == "linregress":
    times, artist_stack, data = tl.linregress_data(conn,type,user,threshold=threshold,interval=interval,stack=stack,centered=centered)
  else:
    times, artist_stack, data = tl.spline_data(conn,type,user,threshold=threshold,interval=interval,stack=stack,centered=centered)
  
  if "-random-colors" in sys.argv:
    color_stack = mtzcolor.colorize_random(artist_stack)
  else:
    begin_color = (0.0,1.0,1.0)
    end_color = (0.8,1.0,1.0)
    color_stack = mtzcolor.colorize_hsv_trans(artist_stack,begin_color,end_color)

  options["color_stack"] = color_stack
  options["artist_stack"] = artist_stack
  options["times"] = times
  options["user"] = user
  options["output_image"] = image
  options["no_display"] = "-no-display" in sys.argv
  options["fill"] = "-no-fill" not in sys.argv
  options["hide_x_axis"] = "-hide-x-axis" in sys.argv
  
  display_graph(data,options)
